#include "Painting.h"
#include "Utility.h"
#include "FitMethods/Utility.h"
using namespace std;
using namespace RooFit;
using namespace Paint;
using namespace Utility;

void Painting::draw(AnalysisConfig config, FitInput fitinput, FitResult fitresult) {
    RooRealVar * isol = new RooRealVar("isol", "E_{T}^{cone40}/E_{T}", -0.2, 1.5);
    gc->add(isol);
    RooPlot* frame = isol->frame();
    RooDataHist * rooDataHist = new RooDataHist("rooDataHist", "rooDataHist", *isol, fitinput.getData());
    gc->add(rooDataHist);
    rooDataHist->plotOn(frame, DataError(RooAbsData::SumW2));
    vector <RooDataHist *> backgroundDataHists = createRooDataHistVector("background", fitinput.getBackgroundTemplates(), *isol, gc);
    vector <RooHistPdf *> backgroundPdfs = createRooHistPdfVector("background", backgroundDataHists, *isol, gc);
    vector <RooDataHist *> signalDataHists = createRooDataHistVector("signal", fitinput.getSignalTemplates(), *isol, gc);
    vector <RooHistPdf *> signalPdfs = createRooHistPdfVector("signal", signalDataHists, *isol, gc);
    for (int i = 0; i < (int) backgroundPdfs.size(); i++) {
        backgroundPdfs[i]->plotOn(frame, DataError(RooAbsData::SumW2), LineStyle(kDashed),
                LineColor(kGreen), Normalization(fitresult.getBackground()[i]->getVal(), RooAbsReal::NumEvent));
    }
    for (int i = 0; i < (int) signalPdfs.size(); i++) {
        signalPdfs[i]->plotOn(frame, DataError(RooAbsData::SumW2), LineStyle(kDotted),
                LineColor(kRed), Normalization(fitresult.getSignal()[i]->getVal(), RooAbsReal::NumEvent));
    }
    //WARNING: Painting's behavior should not depend on what fit method is being used
    if (fitinput.getBackgroundTemplates().size() == fitresult.getBackground().size() &&
            fitinput.getSignalTemplates().size() == fitresult.getSignal().size()) {
        RooAddPdf * resultPdf = createRooPdfAsSum("model", signalPdfs, backgroundPdfs, fitresult.getSignal(), fitresult.getBackground(), gc);
        resultPdf->plotOn(frame, DataError(RooAbsData::SumW2), LineColor(kBlue));
    }
    frame->SetName(fitresult.getName().c_str());

    TPaveText *stats = new TPaveText(0.5, 0.8, 0.95, 0.95, "brNDC");
    stats->SetTextAlign(10);
    for (int i = 0; i < (int) fitresult.getBackground().size(); i++) Utility::addToStats(*stats, *fitresult.getBackground()[i]);
    for (int i = 0; i < (int) fitresult.getSignal().size(); i++) Utility::addToStats(*stats, *fitresult.getSignal()[i]);
    stats->AddText(TString::Format("#chi^{2}/ndof = %5.2f",fitresult.getChiSquare()));

    frame->addObject(stats);
    IOmanager& IOman = IOmanager::instance();
    if (!config.getIsSaveHistograms()) return;

    string dir1 = "DATAPROBE_" + boost::lexical_cast<string > (config.dataProbe.number)
            + "_DATASELECTED_" + boost::lexical_cast<string > (config.dataSelected.number) + "_";
    string dir2 = "ETA_" + boost::lexical_cast<string > (IOman.ETA_VALUE[config.getETAMin()])
            + "_" + boost::lexical_cast<string > (IOman.ETA_VALUE[config.getETAMax() + 1]) + "_";
    string dir3 = "probe_";
    for (int i = 0; i < (int) config.probeSigTempl.size(); i++) dir3 += "SIGPROBE" + boost::lexical_cast<string > (i) + "_" + boost::lexical_cast<string > (config.probeSigTempl[i].number) + "_";
    for (int i = 0; i < (int) config.probeBgdTempl.size(); i++) dir3 += "BGDPROBE" + boost::lexical_cast<string > (i) + "_" + boost::lexical_cast<string > (config.probeBgdTempl[i].number) + "_";
    string dir4 = config.getTag()+"_";
    for (int i = 0; i < (int) config.selectedSigTempl.size(); i++) dir4 += "SIGSELECTED" + boost::lexical_cast<string > (i) + "_" + boost::lexical_cast<string > (config.selectedSigTempl[i].number) + "_";
    for (int i = 0; i < (int) config.selectedBgdTempl.size(); i++) dir4 += "BGDSELECTED" + boost::lexical_cast<string > (i) + "_" + boost::lexical_cast<string > (config.selectedBgdTempl[i].number) + "_";
    string dir5 = "ET_" + boost::lexical_cast<string > (config.getETMin())
            + "_" + boost::lexical_cast<string > (config.getETMax()) + "_";
    string dir6 = "ISOLTHR_" + boost::lexical_cast<string > (config.getIsolThr());

    IOman.savePlot(frame, true, dir1, dir2, dir3, dir4, dir5, dir6);
    delete frame;
}
